package query_understanding.query_understand.tree_progress;

import com.ambiverse.api.model.Entity;
import com.esotericsoftware.kryo.util.ObjectMap;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import edu.stanford.nlp.ling.IndexedWord;
import edu.stanford.nlp.semgraph.SemanticGraph;
import edu.stanford.nlp.semgraph.SemanticGraphCoreAnnotations;
import edu.stanford.nlp.util.CoreMap;
import knowlwdge_graph.knowledge_etl.yago_etl.SchemaDefinition;
import query_understanding.query_language.doc_entity_tree.EntityQueryTree;
import query_understanding.query_language.doc_entity_tree.QueryEdge;
import query_understanding.query_language.doc_entity_tree.QueryNode;
import query_understanding.query_language.doc_entity_tree.QueryNodeTokenMap;
import semantic_compute.entity_linking.schema.EntityLinkingResult;

import java.util.*;

/**
 * Created by julianzliu on 4/23/2017.
 */
public class EntityQueryProgress {
    /*
    一个句子对应一棵查询树
     */

    /*
    1）保存查询解析过程中，需要在不同阶段之间传递的数据结构
    2) 记录查询解析情况，比如是否文本的所有成分都已经被处理
     */


    public String originalQuery;

    public CoreMap sentence;

    public SemanticGraph dependencies;

    public EntityQueryTree initTree;

    // 预处理阶段建立隐射，  实体 -> 对应的词组
    public Map<EntityLinkingResult.LinkedEntity, QueryNodeTokenMap> entityTokenMap;

    public Map<IndexedWord, EntityLinkingResult.LinkedEntity> wordEntityMap;


    public Map<String, IndexedWord> conceptTokenMap;

    public Map<IndexedWord, String> wordConceptMap;


    public PriorityQueue<EntityQueryTree> treePriorityQueue;

    public PriorityQueue<EntityQueryTree> finishedPriorityQueue;


    private ObjectMapper mapper;



    public EntityQueryProgress(){

        this.finishedPriorityQueue = genEmptyPriorityQueue();

        this.wordEntityMap = new HashMap<>();

        this.wordConceptMap = new HashMap<>();

        this.mapper = new ObjectMapper();

    }

    public SemanticGraph getDependencies(){
        return this.sentence.get(SemanticGraphCoreAnnotations.BasicDependenciesAnnotation.class);
    }





    public void extendFinishedQueue(){
        for(EntityQueryTree tree : treePriorityQueue){
            if(tree.generateFinished()){
                if(tree.growthed == false)
                    treePriorityQueue.remove( tree );
                else{
                    EntityQueryTree new_tree = tree.deepCopyTree();
                    finishedPriorityQueue.add( new_tree );
                    tree.growthed = false; //标记
                }
            }
        }
    }







    public void addEntityTokenMap(Map<EntityLinkingResult.LinkedEntity, QueryNodeTokenMap> entityTokenMap){
        this.entityTokenMap = entityTokenMap;
        for(EntityLinkingResult.LinkedEntity mention : entityTokenMap.keySet() ){
            QueryNodeTokenMap nodeTokenMap = entityTokenMap.get( mention );
            this.wordEntityMap.put( nodeTokenMap.mainToken, mention );
            for(IndexedWord word : nodeTokenMap.tokenSet){
                this.wordEntityMap.put( word, mention);
            }
        }
    }

    public void addConceptTokenMap( Map<String, IndexedWord> conceptTokenMap){
        this.conceptTokenMap = conceptTokenMap;
        for(String concept : conceptTokenMap.keySet() ){
            this.wordConceptMap.put(conceptTokenMap.get(concept), concept);
        }
    }


    /**************************
     *
     */
    public JsonNode genCanditateTrees(){
        ObjectNode canditateTrees = this.mapper.createObjectNode();

        canditateTrees.put("query", this.originalQuery);

        ArrayNode unFinishedTrees = this.mapper.createArrayNode();
        canditateTrees.put("unFinishedTrees", unFinishedTrees);

        for(EntityQueryTree tree: this.treePriorityQueue ){
            ObjectNode treeNode = this.mapper.createObjectNode();
            treeNode.put("score", tree.treeScore );
            treeNode.put("tree", tree.genDirectedJsonTree( tree.topicEntity ));
            unFinishedTrees.add( treeNode );
        }

        ArrayNode genFinishedTrees = this.mapper.createArrayNode();
        canditateTrees.put("genFinishedTrees", genFinishedTrees);

        for(EntityQueryTree tree: this.finishedPriorityQueue ){
            ObjectNode treeNode = this.mapper.createObjectNode();
            treeNode.put("score", tree.treeScore );
            treeNode.put("tree", tree.genDirectedJsonTree( tree.topicEntity ));
            genFinishedTrees.add( treeNode );
        }

        return canditateTrees;
    }


    /***************************
     *
     */
    public void printAllTress(){
        System.out.println("[original query = ] " + this.originalQuery + "\n\n");
        System.out.println("finishedQueue = " + this.finishedPriorityQueue.size() + " | " + " unFinishedQueue = " + this.treePriorityQueue.size());
        for(EntityQueryTree tree: treePriorityQueue){
            System.out.println("------------------------------------------------>");
            QueryNode.NodeType nodeType = tree.topicEntity.nodeType;
            System.out.println("  topic_entity_type=" + nodeType);
            if(nodeType == QueryNode.NodeType.ENTITY){
                System.out.println("  topic_entity_id=" + ((QueryNode.EntityNode)tree.topicEntity).yago_id );
            }
            System.out.println("  tree_score = " + tree.treeScore);
            System.out.print("Unvisited: ");
            for(IndexedWord word : tree.unVisited){
                System.out.print(word + " ");
            }
            System.out.println();
            deepFirstTraversal( tree );
            System.out.println("<------------------------------------------------");
        }
        for(EntityQueryTree tree: finishedPriorityQueue){
            System.out.println("------------------------------------------------>");
            QueryNode.NodeType nodeType = tree.topicEntity.nodeType;
            System.out.println("  topic_entity_type=" + nodeType);
            if(nodeType == QueryNode.NodeType.ENTITY){
                System.out.println("  topic_entity_id=" + ((QueryNode.EntityNode)tree.topicEntity).yago_id );
            }
            System.out.println("  tree_score = " + tree.treeScore);
            System.out.println();
            deepFirstTraversal( tree );
            System.out.println("<------------------------------------------------");
        }
    }

    public static void deepFirstTraversal(EntityQueryTree tree){
        //System.out.println("###########################################");
        //Set<QueryNode.BaseNode> visitedNodes = new HashSet<>();
        //Set<QueryEdge.BaseEdge> visitedEdges = new HashSet<>();
        QueryNode.BaseNode root = tree.topicEntity;
        deepFirstTraversalFromNode(tree, root);
        /*
        if(root==null){
            System.out.println("This is An Empty Tree!");
            return;
        }
        String indent = "  ";
        recurrentNode(root,indent,visitedNodes,visitedEdges);
        System.out.println("###########################################");
        */
    }

    public static void deepFirstTraversalFromNode(EntityQueryTree tree, QueryNode.BaseNode node){
        System.out.println("[tree.nodeInvertMap]");
        for(IndexedWord word: tree.nodeInvertMap.keySet()){
            System.out.println("[word] " + word);
        }
        System.out.println("###########################################");
        Set<QueryNode.BaseNode> visitedNodes = new HashSet<>();
        Set<QueryEdge.BaseEdge> visitedEdges = new HashSet<>();
        QueryNode.BaseNode root = node;
        if(root==null){
            System.out.println("This is An Empty Tree!");
            return;
        }
        String indent = "  ";
        recurrentNode(root,indent,visitedNodes,visitedEdges);
        System.out.println("###########################################");
    }

    private static void recurrentNode(QueryNode.BaseNode root, String indent, Set<QueryNode.BaseNode> visitedNodes, Set<QueryEdge.BaseEdge> visitedEdges){
        String printStr = "[" + root.nodeType.toString() + "]";
        if(root.nodeType == QueryNode.NodeType.ENTITY )
            printStr += " " + ((QueryNode.EntityNode)root).yago_id;
        if( root.nodeType == QueryNode.NodeType.CONCEPT)
            printStr += " " + ((QueryNode.ConceptNode)root).yago_id;
        printStr += " hash: " + root.hashCode();
        System.out.println(indent + printStr);
        visitedNodes.add( root );
        for(QueryEdge.BaseEdge edge: root.edgesSet ){
            if(visitedEdges.contains( edge )) continue;;
            recurrentEdge( edge, indent + "  ", visitedEdges, visitedNodes);
        }
    }

    private static void recurrentEdge(QueryEdge.BaseEdge edge, String indent, Set<QueryEdge.BaseEdge> visitedEdges, Set<QueryNode.BaseNode> visitedNodes){
        String printStr = "<" + edge.edgeType.toString() + ">";
        if(edge.edgeType == QueryEdge.EdgeType.YAGO_EDGE )
            printStr += " " + ((QueryEdge.YagoEdge)edge).yago_label;
        if(edge.edgeType == QueryEdge.EdgeType.OCCUR_EDGE)
            printStr += " " + ((QueryEdge.OccurEdge)edge).text;
        printStr += " hash: " + edge.hashCode();
        System.out.println( indent + printStr );
        visitedEdges.add( edge );
        for(QueryNode.BaseNode root: edge.nodesSet ){
            if(visitedNodes.contains(root)) continue;
            recurrentNode(root, indent + " ", visitedNodes, visitedEdges);
        }
    }


    /***************************************
     *
     * @return
     */
    public static PriorityQueue<EntityQueryTree> genEmptyPriorityQueue(){
        Comparator<EntityQueryTree> entityQueryTreeComparator = new TreeScoreComparator();
        PriorityQueue<EntityQueryTree> queryTrees = new PriorityQueue<>(entityQueryTreeComparator);
        return queryTrees;
    }

    public static class TreeScoreComparator implements Comparator<EntityQueryTree> {
        @Override
        public int compare(EntityQueryTree x, EntityQueryTree y){
            if( x.treeScore > y.treeScore) return -1;
            else if(x.treeScore < y.treeScore) return 1;
            return 0;
        }
    }
}
